import pandas as pd
import os
import tqdm
import networkx as nx
import mxnet as mx
root = "/data1/zyx/yks/dataset/retail/wangbo_stage2/ub/"
res_fpn = pd.read_csv(root + "res_ori_fpn.csv")
res_fpn=res_fpn[res_fpn["confindence"]>=0.4]
print(res_fpn.shape)

res_fpn_dcn = pd.read_csv(root +"res_ori_fpn_dcn.csv")
res_fpn_dcn=res_fpn_dcn[res_fpn_dcn["confindence"]>=0.4]
print(res_fpn_dcn.shape)

res_rfcn_dcn = pd.read_csv(root +"res_ori_rfcn_dcn.csv")
res_rfcn_dcn=res_rfcn_dcn[res_rfcn_dcn["confindence"]>=0.4]
print(res_rfcn_dcn.shape)

# res_faster_rcnn = pd.read_csv(root +"faster_rcnn.csv")
# res_faster_rcnn=res_faster_rcnn[res_faster_rcnn["confindence"]>=0.4]
# res_faster_rcnn = res_faster_rcnn[["filename","minx","miny","maxx","maxy","confindence","image_id"]]
# print(res_faster_rcnn.shape)


res_kohill = pd.read_csv(root +"kohill_res.csv",sep=',')
res_kohill=res_kohill[res_kohill["confindence"]>=0.8]
res_kohill_deal = res_kohill[["filename","minx","miny","maxx","maxy","confindence","image_id"]]
print(res_kohill_deal.shape)

res_retinanet = pd.read_csv(root +"retinanet_res.csv",sep=',')
res_retinanet=res_retinanet[res_retinanet["confindence"]>=0.5]
res_retinanet_deal = res_retinanet[["filename","minx","miny","maxx","maxy","confindence","image_id"]]
print(res_retinanet_deal.shape)


res = pd.concat([res_fpn,res_fpn_dcn,res_rfcn_dcn,res_kohill_deal,res_retinanet_deal])
print(res.shape[0])
print(res.head(10))
import numpy as np
final =pd.DataFrame()
iou_type = "kohill"
def is_connected(box0,box1):
    if iou_type =="wangbo":
        x11,x12,y11,y12 = box0[:4]
        x21,x22,y21,y22 = box1[:4]
        sum1 = (x12 - x11) * (y12 - y11)
        sum2 = (x22 - x21) * (y22 - y21)
        sum = min(sum1, sum2)
        area2 = 0
        if min(y22,y12) > max(y11,y21) and min(x12,x22) > max(x11,x21):
            area2 =  (min(y22,y12) - max(y11,y21)) * (min(x12,x22) - max(x11,x21))
        else:
            io = mx.ndarray.contrib.box_iou(mx.nd.array([box0[:4]]), mx.nd.array([box1[:4]]))
            assert io.squeeze().asscalar() < 0.001,[io.squeeze().asscalar(),x21,x22,y21,y22,x11,x12,y11,y12]
        return float(area2) / float(sum) > 0.9
    else:
        io = mx.ndarray.contrib.box_iou(mx.nd.array([box0[:4]]),mx.nd.array([box1[:4]]))
        return io.squeeze().asscalar() > .75
def union_mean(boxes):
    graph = nx.Graph()
    boxes_r = []
    for i in range(boxes.shape[0]):
        for j in range(i + 1, boxes.shape[0]):
            if is_connected(boxes[i],boxes[j]):
                graph.add_edge(i,j)
    nc = nx.connected_components(graph)
    for n in nc:
        if len(n) >= 3:
            boxes_connected = boxes[list(n)]
            boxes_connected = np.mean(boxes_connected[:,:5], axis=0)
#            boxes_connected[:, 4] = np.max(boxes_connected[:, 4], axis=0)
            boxes_r.append([boxes_connected])
    if len(boxes_r) > 0:
        boxes_r = np.concatenate(boxes_r,axis=0)
    boxes_r = np.array(boxes_r)
    return boxes_r
with open(root+"res_combine_kohill_networkx_%s_iou.csv"%(iou_type,),"wt") as f:

    for fname in tqdm.tqdm((res['filename'].unique())):
        dealtrain= res[res['filename'] == fname]
        image_id = list(dealtrain.loc[:, "image_id"])[0]
        bbox_onefile = np.array(dealtrain.loc[:,('minx','miny','maxx','maxy','confindence')]).astype('f')
        bbox_onefile = union_mean(bbox_onefile)
        for bbox in bbox_onefile:
            x0,y0,x1,y1,score = bbox.tolist()
            f.write("%s,%f,%f,%f,%f,%f,%s\r\n"%(fname,x0,y0,x1,y1,score, image_id))
# final.to_csv(root+"res_combine_kohill2.csv",index=False)




